#include "npu.h"

NPU::NPU() {}

NPU::~NPU()
{
    pen.detach_from_chn();
    pen.destroy_region();
    pen_client.detach_from_chn();
    pen_client.destroy_region();
}

bool NPU::Init(const std::string &model_path, bool is_yvu)
{
    g_frame.pool_id = HI_VB_INVALID_POOL_ID;
    ;
    if (memset_s(&g_frame, sizeof(hi_video_frame_info), 0, sizeof(hi_video_frame_info)) != EOK)
    {
        custom_printf("memset_s frame error!!!\n");
        return false;
    }

    if (!detr.Init(model_path))
    {
        return false;
    }
    pen.Init(0, 6, HI_RGN_COVEREX);
    pen_client.Init(1, 7, HI_RGN_COVEREX);
    is_yvu_ = is_yvu;
    return true;
}

/**
 * @brief 调用NPU进行目标检测
 *
 * 从视频帧中检测目标，并将检测结果保存到recv_bbox中。
 *
 * @param recv_bbox 用于存储检测结果的bbox_t结构体引用
 * @param score_thresh 检测结果的置信度阈值
 *
 * @return 如果检测到目标并满足置信度阈值，则返回true；否则返回false
 */
bool NPU::NPUInvoke(bbox_t &recv_bbox, float score_thresh)
{
    memset(&recv_bbox, 0, sizeof(bbox_t));
    hi_s32 ret = hi_mpi_vpss_get_chn_frame(0, 1, &g_frame, 100); //< 获取视频帧
    // custom_printf("color gamut = %d range = %d\n", g_frame.video_frame.color_gamut, g_frame.video_frame.dynamic_range); // OT_COLOR_GAMUT_BT709 limit
    if (ret != HI_SUCCESS) {
        custom_printf("hi_mpi_vpss_get_chn_frame failed! Error Code:%#x!", ret);
        return false;
    }

    std::vector<bbox_t> bboxes;  // 用于存储检测结果的bbox_t结构体数组
    bool recv = detr.detect_uav(&g_frame.video_frame, bboxes, is_yvu_);
    if (hi_mpi_vpss_release_chn_frame(0, 1, &g_frame) != HI_SUCCESS)
    {
        custom_printf("hi_mpi_vpss_release_chn_frame failed! Error Code:%#x!", ret);
    }
    g_frame.pool_id = HI_VB_INVALID_POOL_ID;  // 释放内存
    pen.detach_from_chn(); // 强制解绑旧区域 
    for (auto& bbox : bboxes) {
        if (bbox.class_id == 0 && bbox.score >= score_thresh) {
            std::cout << "target_bbox.x = " << bbox.x << ", target_bbox.y = " << bbox.y << ", target_bbox.w = " << bbox.w << ", target_bbox.h = " << bbox.h << std::endl;
            recv_bbox = bbox; 
            pen.detach_from_chn();  
            pen.set_attr(recv_bbox, 0xFFFF00); 
            pen.attach_to_chn();
            return true;
        }
    }
    return false;
}

bool NPU::NPUInvoke_client(bbox_t &recv_bbox, float score_thresh) 
{
    memset(&recv_bbox, 0, sizeof(bbox_t));
    hi_s32 ret = hi_mpi_vpss_get_chn_frame(0, 1, &g_frame, 100); //< 获取视频帧
    // custom_printf("color gamut = %d range = %d\n", g_frame.video_frame.color_gamut, g_frame.video_frame.dynamic_range); // OT_COLOR_GAMUT_BT709 limit
    if (ret != HI_SUCCESS) {
        custom_printf("hi_mpi_vpss_get_chn_frame failed! Error Code:%#x!", ret);
        return false;
    }

    std::vector<bbox_t> bboxes; // 用于存储检测结果的bbox_t结构体数组
    bool recv = detr.detect_uav_client(&g_frame.video_frame, bboxes, is_yvu_);
    if (hi_mpi_vpss_release_chn_frame(0, 1, &g_frame) != HI_SUCCESS) {
        custom_printf("hi_mpi_vpss_release_chn_frame failed! Error Code:%#x!", ret);
    }
    g_frame.pool_id = HI_VB_INVALID_POOL_ID;  // 释放内存
    pen_client.detach_from_chn(); // 强制解绑旧区域 
    for (auto& bbox : bboxes) {
        if (bbox.class_id == 0 && bbox.score >= score_thresh) {
            std::cout << "target_bbox.x = " << bbox.x << ", target_bbox.y = " << bbox.y << ", target_bbox.w = " << bbox.w << ", target_bbox.h = " << bbox.h << std::endl;
            recv_bbox = bbox; 
            pen_client.detach_from_chn();  
            pen_client.set_attr(recv_bbox, 0xFFFF00); 
            pen_client.attach_to_chn();
            return true;
        }
    }   
    return false;
}